from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts.prompt import PromptTemplate
from langchain.chains import LLMChain

import os

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
# 定义模型
llm = Tongyi()

# 定义两个chain
synopsis_template = """你是一名写作家. 给定游戏的标题,你的工作就是为该标题写一个概要.

标题: {title}
写作家: 这是以上游戏的简介:"""
synopsis_prompt_template = PromptTemplate(input_variables=["title"], template=synopsis_template)
synopsis_chain = LLMChain(llm=llm, prompt=synopsis_prompt_template, output_key="content")

commentator_template = """你是一名评论员. 给定游戏的简介,你的工作就是为该游戏，给出评分（1-10）.

游戏简介概要: {content}
评论员: 这是对于该游戏的评分:"""
commentator_prompt_template = PromptTemplate(input_variables=["content"], template=commentator_template)
commentator_chain = LLMChain(llm=llm, prompt=commentator_prompt_template)

# 顺序调用chain
from langchain.chains import SimpleSequentialChain

overall_chain = SimpleSequentialChain(chains=[synopsis_chain, commentator_chain],
                                      input_key="title",
                                      output_key="comment"
                                      )
# 第一个chain的结果作为第二个chain中{content}的入参
res = overall_chain.invoke({"title": '奥特曼大战怪兽'})
print(res)
